from PySide6.QtWidgets import QApplication,QMainWindow,QPushButton,QVBoxLayout,QWidget,QFileDialog,QLabel,QHBoxLayout
from PySide6.QtGui import QPixmap,QDesktopServices
from PySide6.QtCore import Qt,QUrl
# from hats_detection.video_worker import worker_process_video
# from hats_detection.image_worker import worker_process_image
from hats_detection.model_loader import ModelLoader
from hats_detection.workers import GenericWorker
from hats_detection.video_processor import VideoProcessor
from hats_detection.image_processor import ImageProcessor


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("帽子检测系统")

        #创建QLabel 用于显示处理后的内容
        self.content_label = QLabel(self)
        self.content_label.resize(640,480)

        #创建QLabel用于显示模型的路径
        self.model_path_label = QLabel("模型路径：未加载",self)

        #创建读取模型按钮
        self.load_model_button = QPushButton("加载模型",self)
        self.load_model_button.clicked.connect(self.load_model)

        #创建处理视频按钮
        self.process_video_button = QPushButton("处理视频",self)
        self.process_video_button.clicked.connect(self.process_video)
        self.process_video_button.setEnabled(False)

        #创建处理图片按钮
        self.process_image_button = QPushButton("处理图片",self)
        self.process_image_button.clicked.connect(self.process_image)
        self.process_image_button.setEnabled(False)

        #创建停止按钮
        self.stop_button = QPushButton("停止",self)
        self.stop_button.clicked.connect(self.stop_processing)
        self.stop_button.setEnabled(False)

        #创建布局
        layout = QVBoxLayout()

        layout.addWidget(self.model_path_label)

        #创建一个存按钮的布局
        button_layout = QHBoxLayout()

        #把按钮添加到按钮布局中
        button_layout.addWidget(self.load_model_button)
        button_layout.addWidget(self.process_video_button)
        button_layout.addWidget(self.process_image_button)
        button_layout.addWidget(self.stop_button)

        #把按钮布局添加到布局中
        layout.addLayout(button_layout)

        #把内容标签添加到布局中
        layout.addWidget(self.content_label)
        

        container = QWidget()
        container.setLayout(layout)
        self.setCentralWidget(container)

        #初始化组件
        self.model_loader = ModelLoader()

        self.model_loader.model_loaded.connect(self.on_model_loaded)

        self.video_processor = VideoProcessor()
        self.video_processor.change_pixmap_signal.connect(self.update_image)

        self.image_processor = ImageProcessor()

        self.model = None
        self.load_model_worker = None
        

    def load_model(self):
        model_path,_ = QFileDialog.getOpenFileName(self,"选择模型文件","","*.pt *.pth")
        if model_path:
            #创建了一个GenericWorker 线程任务，用于异步加载模型，防止界面卡顿
            self.load_model_worker = GenericWorker(lambda: self.model_loader.load_model(model_path))
            # self.load_model_worker = GenericWorker(self._load_model_task,model_path)    #一般写法
            # 当线程任务完成后，需要自动删除线程对象以释放资源
            self.load_model_worker.finished.connect(self.load_model_worker.deleteLater)
            #启动线程任务，开始加载模型
            self.load_model_worker.start()

    def on_model_loaded(self,model):
        if model:
            self.model = model
            self.video_processor.set_model(model)
            self.image_processor.set_model(model)
            self.process_video_button.setEnabled(True)
            self.process_image_button.setEnabled(True)
            self.model_path_label.setText(f"模型路径: {self.model_loader.model_path}")
        


    # def _load_model_task(self,model_path):    # 一般写法的函数
    #     return self.model_loader.load_model(model_path)

    def process_video(self):
        video_path,_ = QFileDialog.getOpenFileName(self,"选择视频文件","","*.mp4")
        if video_path and self.model:
            self.video_processor.set_video_path(video_path)
            self.video_processor.start()
            self.stop_button.setEnabled(True)

    def update_image(self,image):
        self.content_label.setPixmap(QPixmap.fromImage(image))
    
    def process_image(self):
        image_path,_ = QFileDialog.getOpenFileName(self,"选择图片文件","","*.jpg")
        if image_path and self.model:
            qt_image = self.image_processor.process_image(image_path)
            if qt_image:
                self.content_label.setPixmap(QPixmap.fromImage(qt_image))
        

    def stop_processing(self):
        self.video_processor.stop()
        self.video_processor.wait()
        
